{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "executionInfo": {
     "elapsed": 350,
     "status": "ok",
     "timestamp": 1649955640910,
     "user": {
      "displayName": "Sam Lu",
      "userId": "15789059763790170725"
     },
     "user_tz": -480
    },
    "id": "BANe9WM-_Ew0"
   },
   "outputs": [],
   "source": [
    "import gymnasium as gym\n",
    "import torch\n",
    "import torch.nn.functional as F\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import rl_utils"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "executionInfo": {
     "elapsed": 2,
     "status": "ok",
     "timestamp": 1649955643798,
     "user": {
      "displayName": "Sam Lu",
      "userId": "15789059763790170725"
     },
     "user_tz": -480
    },
    "id": "n1peNzvo_Ew1"
   },
   "outputs": [],
   "source": [
    "## 构造智能体 agent 的大脑，也就是输入状态，返回该状态下，选择每个动作的概率\n",
    "## 输入是状态的，也就是 (车子center-point的坐标，车子的速度，杆的竖直角度，杆的角速度)\n",
    "## 返回值应该是2 dim\n",
    "class PolicyNet(torch.nn.Module):\n",
    "    def __init__(self, state_dim, hidden_dim, action_dim):\n",
    "        super(PolicyNet, self).__init__()\n",
    "        self.fc1 = torch.nn.Linear(state_dim, hidden_dim)\n",
    "        self.fc2 = torch.nn.Linear(hidden_dim, action_dim)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = F.relu(self.fc1(x))\n",
    "        return F.softmax(self.fc2(x), dim=1)  ## 返回该状态下，选择的动作的概率"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "executionInfo": {
     "elapsed": 2,
     "status": "ok",
     "timestamp": 1649955645289,
     "user": {
      "displayName": "Sam Lu",
      "userId": "15789059763790170725"
     },
     "user_tz": -480
    },
    "id": "KvWZvDoO_Ew1"
   },
   "outputs": [],
   "source": [
    "## 构造智能体 agent 的大脑，也就是输入状态，返回该状态下，每个动作的动作价值\n",
    "## 输入是状态的，也就是 (车子center-point的坐标，车子的速度，杆的竖直角度，杆的角速度)\n",
    "## 返回值应该是2 dim\n",
    "class ValueNet(torch.nn.Module):\n",
    "    def __init__(self, state_dim, hidden_dim):\n",
    "        super(ValueNet, self).__init__()\n",
    "        self.fc1 = torch.nn.Linear(state_dim, hidden_dim)\n",
    "        self.fc2 = torch.nn.Linear(hidden_dim, 1)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = F.relu(self.fc1(x))\n",
    "        return self.fc2(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "executionInfo": {
     "elapsed": 516,
     "status": "ok",
     "timestamp": 1649955673254,
     "user": {
      "displayName": "Sam Lu",
      "userId": "15789059763790170725"
     },
     "user_tz": -480
    },
    "id": "FBnRPG_n_Ew2"
   },
   "outputs": [],
   "source": [
    "## 智能体\n",
    "class ActorCritic:\n",
    "    def __init__(self, state_dim, hidden_dim, action_dim, actor_lr, critic_lr,\n",
    "                 gamma, device):\n",
    "        # 策略网络 Actor\n",
    "        self.actor = PolicyNet(state_dim, hidden_dim, action_dim).to(device)\n",
    "        self.critic = ValueNet(state_dim, hidden_dim).to(device)  # 价值网络  Critic\n",
    "        # 策略网络优化器\n",
    "        self.actor_optimizer = torch.optim.Adam(self.actor.parameters(),\n",
    "                                                lr=actor_lr)\n",
    "        self.critic_optimizer = torch.optim.Adam(self.critic.parameters(),\n",
    "                                                 lr=critic_lr)  # 价值网络优化器\n",
    "        self.gamma = gamma  # 折扣因子\n",
    "        self.device = device\n",
    "\n",
    "    def take_action(self, state):            # 根据动作概率分布随机采样\n",
    "        state = torch.tensor([state], dtype=torch.float).to(self.device)\n",
    "        probs = self.actor(state)       ## 拿到该状态下，每个动作的选择概率\n",
    "        action_dist = torch.distributions.Categorical(probs)    ##   配置 好采样的概率\n",
    "        action = action_dist.sample()        ## 对该状态下，所有的动作采样，采样的概率是probs\n",
    "        return action.item()                 ## 返回依概率采样得到的动作\n",
    "\n",
    "    ## 训练策略网络的，用一条序列来训练\n",
    "    ## 不用遍历了的，可以批量来处理，因不需要求每个状态的回报，价值函数迭代不需要遍历，直接求就可以的呢\n",
    "    def update(self, transition_dict):\n",
    "        ## 拿到这条序列内的 奖励、状态和动作\n",
    "        states = torch.tensor(transition_dict['states'],\n",
    "                              dtype=torch.float).to(self.device)\n",
    "        actions = torch.tensor(transition_dict['actions']).view(-1, 1).to(\n",
    "            self.device)\n",
    "        rewards = torch.tensor(transition_dict['rewards'],\n",
    "                               dtype=torch.float).view(-1, 1).to(self.device)\n",
    "        next_states = torch.tensor(transition_dict['next_states'],\n",
    "                                   dtype=torch.float).to(self.device)\n",
    "        dones = torch.tensor(transition_dict['dones'],\n",
    "                             dtype=torch.float).view(-1, 1).to(self.device)\n",
    "\n",
    "        # 时序差分目标\n",
    "        ## 用下一个状态，critic求出下个状态的动作价值，然后求出当前状态的动作价值\n",
    "        td_target = rewards + self.gamma * self.critic(next_states) * (1 - dones)  ## 真实标签的，truth label，有监督\n",
    "        ## critic使用当前状态，求出当前状态的动作价值，两者的差就是差分\n",
    "        td_delta = td_target - self.critic(states)  # 时序差分误差\n",
    "        log_probs = torch.log(self.actor(states).gather(1, actions))     ## 选择的动作的动作概率，并求 log\n",
    "        ## 策略网络的损失，差分越小越好，-log_probs > 0，td_delta.detach()也就是不用反向求梯度，这的td_delta看作是固定的值\n",
    "        ## -log_probs > 0, 所以越靠近0越好，当q_probably=1时最小，也就是选择的动作概率越大越好, td_delta越来越小，log_probs要越来越大才行\n",
    "        actor_loss = torch.mean(-log_probs * td_delta.detach())          ## 时序差分误差，乘上相应的 log值，就得到策略网络的损失loss\n",
    "        \n",
    "        ## 均方误差损失函数，价值网络critic求出当前状态的动作价值，以及用下一个状态间接求出当前状态的动作价值，MSE求损失loss\n",
    "        ## 价值网络的损失，td_target.detach()不用反向求梯度，所以td_target看作truth label，self.critic(states)看作predict\n",
    "        critic_loss = torch.mean(F.mse_loss(self.critic(states), td_target.detach()))\n",
    "        self.actor_optimizer.zero_grad()     ## 参数的梯度置0\n",
    "        self.critic_optimizer.zero_grad()    ## 参数的梯度置0\n",
    "        actor_loss.backward()  # 计算策略网络的梯度\n",
    "        critic_loss.backward()  # 计算价值网络的梯度\n",
    "        self.actor_optimizer.step()  # 更新策略网络的参数\n",
    "        self.critic_optimizer.step()  # 更新价值网络的参数"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 52404,
     "status": "ok",
     "timestamp": 1649955727483,
     "user": {
      "displayName": "Sam Lu",
      "userId": "15789059763790170725"
     },
     "user_tz": -480
    },
    "id": "bJ5VTKxU_Ew3",
    "outputId": "c70fd2c1-dcc9-4a11-884e-ff463fdb1069"
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Iteration 0:   0%|                                                                                                     | 0/100 [00:00<?, ?it/s]\n"
     ]
    },
    {
     "ename": "ValueError",
     "evalue": "expected sequence of length 4 at dim 2 (got 0)",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mValueError\u001b[0m                                Traceback (most recent call last)",
      "Cell \u001b[1;32mIn[13], line 18\u001b[0m\n\u001b[0;32m     14\u001b[0m action_dim \u001b[38;5;241m=\u001b[39m env\u001b[38;5;241m.\u001b[39maction_space\u001b[38;5;241m.\u001b[39mn\n\u001b[0;32m     15\u001b[0m agent \u001b[38;5;241m=\u001b[39m ActorCritic(state_dim, hidden_dim, action_dim, actor_lr, critic_lr,\n\u001b[0;32m     16\u001b[0m                     gamma, device)\n\u001b[1;32m---> 18\u001b[0m return_list \u001b[38;5;241m=\u001b[39m \u001b[43mrl_utils\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mtrain_on_policy_agent\u001b[49m\u001b[43m(\u001b[49m\u001b[43menv\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43magent\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mnum_episodes\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[1;32m~\\Desktop\\access\\Hands-on-RL\\rl_utils.py:40\u001b[0m, in \u001b[0;36mtrain_on_policy_agent\u001b[1;34m(env, agent, num_episodes)\u001b[0m\n\u001b[0;32m     38\u001b[0m done \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mFalse\u001b[39;00m\n\u001b[0;32m     39\u001b[0m \u001b[38;5;66;03m## 采样一条序列的\u001b[39;00m\n\u001b[1;32m---> 40\u001b[0m \u001b[38;5;28;01mwhile\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m done:\n\u001b[0;32m     41\u001b[0m     action \u001b[38;5;241m=\u001b[39m agent\u001b[38;5;241m.\u001b[39mtake_action(state)    \u001b[38;5;66;03m##  根据状态采取动作的\u001b[39;00m\n\u001b[0;32m     42\u001b[0m     \u001b[38;5;66;03m##  环境执行动作，并反馈下一个状态、动作的奖励、是否完成、步长太长的，info\u001b[39;00m\n",
      "Cell \u001b[1;32mIn[12], line 17\u001b[0m, in \u001b[0;36mActorCritic.take_action\u001b[1;34m(self, state)\u001b[0m\n\u001b[0;32m     16\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mtake_action\u001b[39m(\u001b[38;5;28mself\u001b[39m, state):            \u001b[38;5;66;03m# 根据动作概率分布随机采样\u001b[39;00m\n\u001b[1;32m---> 17\u001b[0m     state \u001b[38;5;241m=\u001b[39m \u001b[43mtorch\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mtensor\u001b[49m\u001b[43m(\u001b[49m\u001b[43m[\u001b[49m\u001b[43mstate\u001b[49m\u001b[43m]\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mdtype\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mtorch\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mfloat\u001b[49m\u001b[43m)\u001b[49m\u001b[38;5;241m.\u001b[39mto(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mdevice)\n\u001b[0;32m     18\u001b[0m     probs \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mactor(state)       \u001b[38;5;66;03m## 拿到该状态下，每个动作的选择概率\u001b[39;00m\n\u001b[0;32m     19\u001b[0m     action_dist \u001b[38;5;241m=\u001b[39m torch\u001b[38;5;241m.\u001b[39mdistributions\u001b[38;5;241m.\u001b[39mCategorical(probs)    \u001b[38;5;66;03m##   配置 好采样的概率\u001b[39;00m\n",
      "\u001b[1;31mValueError\u001b[0m: expected sequence of length 4 at dim 2 (got 0)"
     ]
    }
   ],
   "source": [
    "actor_lr = 1e-3\n",
    "critic_lr = 1e-2\n",
    "num_episodes = 1000\n",
    "hidden_dim = 128\n",
    "gamma = 0.98\n",
    "device = torch.device(\"cuda\") if torch.cuda.is_available() else torch.device(\n",
    "    \"cpu\")\n",
    "\n",
    "env_name = 'CartPole-v1'\n",
    "env = gym.make(env_name)\n",
    "_ = env.reset(seed=0)\n",
    "torch.manual_seed(0)\n",
    "state_dim = env.observation_space.shape[0]\n",
    "action_dim = env.action_space.n\n",
    "agent = ActorCritic(state_dim, hidden_dim, action_dim, actor_lr, critic_lr,\n",
    "                    gamma, device)\n",
    "\n",
    "return_list = rl_utils.train_on_policy_agent(env, agent, num_episodes)\n",
    "\n",
    "# Iteration 0: 100%|██████████| 100/100 [00:00<00:00, 184.32it/s, episode=100,\n",
    "# return=21.100]\n",
    "# Iteration 1: 100%|██████████| 100/100 [00:01<00:00, 98.31it/s, episode=200,\n",
    "# return=72.800]\n",
    "# Iteration 2: 100%|██████████| 100/100 [00:01<00:00, 58.72it/s, episode=300,\n",
    "# return=109.300]\n",
    "# Iteration 3: 100%|██████████| 100/100 [00:04<00:00, 23.14it/s, episode=400,\n",
    "# return=163.000]\n",
    "# Iteration 4: 100%|██████████| 100/100 [00:08<00:00, 11.78it/s, episode=500,\n",
    "# return=193.600]\n",
    "# Iteration 5: 100%|██████████| 100/100 [00:08<00:00, 11.23it/s, episode=600,\n",
    "# return=195.900]\n",
    "# Iteration 6: 100%|██████████| 100/100 [00:08<00:00, 11.55it/s, episode=700,\n",
    "# return=199.100]\n",
    "# Iteration 7: 100%|██████████| 100/100 [00:09<00:00, 10.75it/s, episode=800,\n",
    "# return=186.900]\n",
    "# Iteration 8: 100%|██████████| 100/100 [00:08<00:00, 11.73it/s, episode=900,\n",
    "# return=200.000]\n",
    "# Iteration 9: 100%|██████████| 100/100 [00:08<00:00, 12.05it/s, episode=1000,\n",
    "# return=200.000]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 573
    },
    "executionInfo": {
     "elapsed": 6,
     "status": "ok",
     "timestamp": 1649955727484,
     "user": {
      "displayName": "Sam Lu",
      "userId": "15789059763790170725"
     },
     "user_tz": -480
    },
    "id": "3mXhNUJU_Ew4",
    "outputId": "f2a53ebf-58ec-499c-c740-36ed419941c8"
   },
   "outputs": [],
   "source": [
    "episodes_list = list(range(len(return_list)))\n",
    "plt.plot(episodes_list, return_list)\n",
    "plt.xlabel('Episodes')\n",
    "plt.ylabel('Returns')\n",
    "plt.title('Actor-Critic on {}'.format(env_name))\n",
    "plt.show()\n",
    "\n",
    "mv_return = rl_utils.moving_average(return_list, 9)\n",
    "plt.plot(episodes_list, mv_return)\n",
    "plt.xlabel('Episodes')\n",
    "plt.ylabel('Returns')\n",
    "plt.title('Actor-Critic on {}'.format(env_name))\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "colab": {
   "collapsed_sections": [],
   "name": "第10章-Actor-Critic算法.ipynb",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
